from strategy.btStg.Sma import Sma
from strategy.btStg.SmaCross import SmaCross
from tools import dataTool, dbTool
import backtrader as bt


def multiTest(stg):
    codes = dataTool.getCalcTradeCode()
    for code in codes:  # 测试多只股票
        data = dataTool.getBtKdata(code, Config.start, Config.end)
        paramList = stg.getParamList()
        for param in paramList:
            try:
                test(stg, param, code, data)
            except BaseException as e:
                print(e)


def oneTest(stg, param, code):
    data = dataTool.getBtKdata(code, Config.start, Config.end)
    test(stg, param, code, data)


def test(stg, param, code, data):
    # 基本配置
    sizer = Config.sizer
    input = Config.input
    # 创建大脑
    cerebro = bt.Cerebro()
    # 添加策略
    cerebro.addstrategy(stg, param=param)
    cerebro.broker.setcash(input)  # 设置初始资金

    # cerebro.broker.setcommission(0.001)  # 设置手续费
    cerebro.adddata(data)
    cerebro.addsizer(bt.sizers.AllInSizer)
    # cerebro.addsizer(sizer)
    # 运行
    cerebro.run()
    output = round(cerebro.broker.getvalue())
    profitRate = round((output - input) / input, 3)
    sizerName = sizer.__name__
    stgName = stg.__name__
    result = {'stg': stgName, 'param': param, 'code': code,
              'input': input, 'output': output, 'sizer': sizerName, 'profit_rate': profitRate}
    # dbTool.saveOne('result_bt', result)
    print(result)
    cerebro.plot()
    # exit()


class Config:
    start = '20160301'
    end = '20171231'
    input = 100000
    sizer = bt.sizers.AllInSizer


# multiTest(SmaCross)
# oneTest(SmaCross, {'periodFast': 10, 'periodSlow': 50}, '600855.SH')
oneTest(Sma, {'period': 10}, '000009.SZ')
